{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "414d6ab3-d602-4322-bfc4-861380192563",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from params import *\n",
    "from mass_assign import CIC_3D\n",
    "from astropy.convolution import convolve\n",
    "from photutils.segmentation import make_2dgaussian_kernel, detect_threshold,  detect_sources\n",
    "from photutils.isophote import Ellipse, EllipseGeometry, build_ellipse_model\n",
    "from photutils import EllipticalAperture"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c39b1866-8e3e-447e-b9ae-44446e766e40",
   "metadata": {},
   "source": [
    "## Figure 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7e900f2-9816-4cb8-ba27-bf4e3d5fa32b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "snapshot = np.arange(5, 125, 7)\n",
    "plt.figure(figsize=(9, len(snapshot)*1.56/2))\n",
    "for n,i in enumerate(snapshot):\n",
    "    part = np.load(\"./data/\" + str(i) + \".npy\")\n",
    "    \n",
    "    dens = CIC_3D(part)\n",
    "    kernel = make_2dgaussian_kernel(fwhm=5, size=15)\n",
    "    \n",
    "    plt.subplot(len(snapshot)//2,6,3*n+1)\n",
    "    intens = convolve(np.sum(dens, axis=0) + np.random.randn(512, 512)*.1, kernel)\n",
    "    intens[intens < 0] = 0\n",
    "    plt.imshow(np.log(intens + .01), cmap=\"gray\", origin=\"lower\")\n",
    "    plt.xticks([]); plt.yticks([]); \n",
    "    plt.text(5, 500, str(round(i*0.01, 2))+\" Gyr\", c=\"w\", ha=\"left\", va=\"top\", fontsize=\"small\")\n",
    "    if n < 2: plt.title(\"xy\")\n",
    "    plt.subplot(len(snapshot)//2,6,3*n+2)\n",
    "    intens = convolve(np.sum(dens, axis=1) + np.random.randn(512, 512)*.1, kernel)\n",
    "    intens[intens < 0] = 0\n",
    "    plt.imshow(np.log(intens + .01), cmap=\"gray\", origin=\"lower\")\n",
    "    plt.xticks([]); plt.yticks([]) \n",
    "    if n < 2: plt.title(\"xz\")\n",
    "    plt.subplot(len(snapshot)//2,6,3*n+3)\n",
    "    intens = convolve(np.sum(dens, axis=2) + np.random.randn(512, 512)*.1, kernel)\n",
    "    intens[intens < 0] = 0\n",
    "    plt.imshow(np.log(intens + .01), cmap=\"gray\", origin=\"lower\")\n",
    "    plt.xticks([]); plt.yticks([]) \n",
    "    if n < 2: plt.title(\"yz\")\n",
    "    \n",
    "plt.subplots_adjust(wspace=0.02, hspace=0.01)\n",
    "plt.savefig(\"merging.pdf\", dpi=150)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7c1d88f-58d3-463d-bcd8-d11b422a8d17",
   "metadata": {},
   "source": [
    "## Detect tidal features and measure f_tidal (Huang & Fan 2022)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b72353fb-00ce-45d3-8476-e92d6ab06175",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "i = 45\n",
    "part = np.load(\"./data/\" + str(i) + \".npy\")\n",
    "    \n",
    "dens = CIC_3D(part)\n",
    "kernel = make_2dgaussian_kernel(fwhm=15, size=61)\n",
    "\n",
    "intens = convolve(np.sum(dens, axis=0) + np.random.randn(512, 512)*.1, kernel)\n",
    "plt.imshow(np.log(intens +.05))\n",
    "plt.show()\n",
    "\n",
    "img_masked = intens\n",
    "DIM = 512\n",
    "\n",
    "params = {}\n",
    "params['e0'] = 0.2\n",
    "params['pa0'] = 0\n",
    "params['sma0'] = 20\n",
    "params['xcen0'] = 270\n",
    "params['ycen0'] = 300\n",
    "params['pixscl'] = 1\n",
    "\n",
    "# find center\n",
    "geometry = EllipseGeometry(params['xcen0'], params['ycen0'], params['sma0'], params['e0'], params['pa0'])\n",
    "ellipse = Ellipse(img_masked, geometry)\n",
    "iso_free = ellipse.fit_image(fix_center=False, fix_pa=False, fix_eps=False, \n",
    "                             minsma=.2, maxsma=DIM/2*1.35, step=0.15, maxgerr=10)\n",
    "\n",
    "kk = iso_free.sma*params['pixscl'] < 10\n",
    "xcen = np.mean(iso_free.x0[kk])\n",
    "ycen = np.mean(iso_free.y0[kk])\n",
    "params.update({'xcen':xcen, 'ycen':ycen})\n",
    "print(xcen, ycen)\n",
    "\n",
    "# fixed center\n",
    "geometry = EllipseGeometry(params['xcen'], params['ycen'], params['sma0'], params['e0'], params['pa0'])\n",
    "ellipse = Ellipse(img_masked, geometry)\n",
    "\n",
    "iso_fixCen = ellipse.fit_image(fix_center=True, fix_pa=False, fix_eps=False, \n",
    "                               minsma=.2, maxsma=DIM/2*1.3, step=0.1, maxgerr=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "307cbb0f-e530-4670-9119-000c60b483fb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_image = build_ellipse_model(img_masked.shape, iso_fixCen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47781b1b-51ea-45b5-b55a-a3b0f5b8700f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rms = detect_threshold(intens, nsigma=3)\n",
    "residual = intens - model_image\n",
    "mask = (intens/model_image <= 1.15) | (intens < rms)\n",
    "residual[mask] = 0\n",
    "segm = detect_sources(residual, threshold=0, npixels=200)\n",
    "mask = mask | (segm.data<1)\n",
    "residual[mask] = 0\n",
    "\n",
    "plt.imshow(np.log(intens + .05), origin=\"lower\", cmap=\"gray\")\n",
    "plt.imshow(residual, cmap=\"jet\", origin=\"lower\", alpha=.3, vmax=.01)\n",
    "\n",
    "ftidal = round(np.sum(residual) / (np.sum(model_image) + np.sum(residual)), 6)\n",
    "print(ftidal)\n",
    "\n",
    "f_tidal = pd.read_csv(\"f_tidal.csv\")\n",
    "\n",
    "f_tidal.loc[np.where(f_tidal[\"i\"] == i)[0][0], \"ftidal\"] = ftidal\n",
    "f_tidal.to_csv(\"f_tidal.csv\", index=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32ee477b-1303-4595-a190-01a13226ad4c",
   "metadata": {},
   "source": [
    "## Figure 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d89a13f-9a17-4cdc-9354-c64ce6c1fa78",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.rcParams[\"font.size\"] = 14\n",
    "f_tidal = pd.read_csv(\"f_tidal.csv\")\n",
    "plt.scatter((f_tidal[\"i\"][45:84]-45)*.01, f_tidal[\"ftidal\"][45:84], marker=\"x\", c=\"r\")\n",
    "plt.ylim(-0.005, 0.135)\n",
    "\n",
    "expx = np.arange(0, 0.39, 0.005)\n",
    "a, b, c = 0.126, -11.9, 0.009\n",
    "plt.plot(expx, a * np.exp(b * expx) + c, c=\"k\")\n",
    "plt.xlabel(\"time since merging [Gyr]\")\n",
    "plt.ylabel(\"$f_\\mathrm{tidal}$\")\n",
    "plt.legend((\"data\", \"exp. fit\"))\n",
    "plt.savefig(\"f_tidal.pdf\", dpi=250)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
